from numpy import sum,where,array
from dolfin import Function,dx,assemble,errornorm,interpolate,pi,MeshFunction,DirichletBC,VectorFunctionSpace,Expression,DomainBoundary,refine,\
                   FunctionSpace, project, RectangleMesh,plot,File
from inflowboundary import *

class ProblemBase:
    'Base class for the problems.'
    def __init__(self,options):
        self.N = options['N']                                                 # spatial resolution
        self.domain = None                                                    #! domain of the problem
        self.phi_ = None                                                      #! initial conditions
        self.weakBC = False                                                   #! for CG - should the surface term be neglected
        self.ibcTuple = ()                                                    #! (obj. derived from Subdomain,Dirichlet values,ds(0))
        self.t = 0                                                            #! start time
        self.T = None                                                         #! end time
        self.vStrings = None                                                  #! strings for velocity compononets    
        self.vIsTimeDependent = None                                          #! tag for compilation of velocity as an Expression
        self.resultsDir = options['results-path']                             # parent directory for results
        self.problemDir = self.__class__.__name__.lower()                     # 1st sub directory for results
        self.hasExactSolution = None                                          #! for postprocessing, determines number of used variables
        self.activeVariables = None                                           #! variables that are updated
        self.outFileT = 'data-trans.txt'                                      # name of file for output
        self.outFileA = 'data-avg.txt'                                        # for transianets and averages
        self.ibcExpr = None                                                  #! Dolfin expression that determines inflow boundary values
        self.L = None                                                        # length of the interface

    def update(self,phi,t,dt):
        '''Update lists of recorded variables.'''
        self.ts.append(t)
        self.dts.append(dt)

        V = phi.function_space()
        av, ave = self.get_volume(phi,V)
        self.avs.append(av)
        self.aves.append(ave)

        if self.hasExactSolution:
            phiExact = self.exact_solution(t)
            ev, eve = self.get_volume(phiExact,V)
            self.evs.append(ev)
            self.eves.append(eve)
           
            p = V.ufl_element().degree()
            e = errornorm(phiExact,phi,norm_type='L2',degree_rise=p+2)  
            self.es.append(e)

            File('numerical.pvd') << phi, int(t/dt)
            File('exact.pvd') << phiExact, int(t/dt)

            enright = self.enright_measure(phiExact,phi,V)
            self.enrights.append(enright)

    def enright_measure(self,phi_exact,phi,space):
        '''Accuracy measure from enright thesis.'''
        domain = RectangleMesh(self.LL[0],self.LL[1],self.UR[0],self.UR[1],100,100)   # 1000, too much 
        family = space.ufl_element().family()
        degree = space.ufl_element().degree()
        V = FunctionSpace(domain,family,degree)
        
        dense_phi = interpolate(phi,V)                #TODO consider interpolate here
        inside_phi = Function(V)
        inside_phi.vector()[:] = where(dense_phi.vector().array() < 0, 1, 0)

        dense_phi_exact = interpolate(phi_exact,V)
        inside_phi_exact = Function(V)
        inside_phi_exact.vector()[:] = where(dense_phi_exact.vector().array() < 0, 1, 0)

        volumeForm = abs(inside_phi_exact-inside_phi)*dx
        enright = assemble(volumeForm)
        enright = enright/self.L
        
        return enright


    def get_volume(self,dForE,dSpace):
        '''Calculate current volume (and error) from the solution (and exact solution).'''
        # project solution to denser mesh to compute volume
        domain = RectangleMesh(self.LL[0],self.LL[1],self.UR[0],self.UR[1],100,100)   # 1000, too much 
        family = dSpace.ufl_element().family()
        degree = dSpace.ufl_element().degree()
        V = FunctionSpace(domain,family,degree)
        dF = interpolate(dForE,V)                #TODO consider interpolate here

        inside = Function(V)
        inside.vector()[:] = where(dF.vector().array() < 0, 1, 0)

        volumeForm = inside*dx
        exact = self.exact_volume()
        v = assemble(volumeForm)
        ve = 100*abs(v - exact)/exact

        return v, ve
       
    def register_variables(self):
        '''Define attributes that hold results.'''
        self.ts = []
        self.dts = []
        self.avs = []
        self.aves = []
        self.activeVariables = [0,1,2,3]
        if self.hasExactSolution == True:
            self.evs = []
            self.eves = []
            self.es = []
            self.enrights = []
            self.activeVariables += [4,5,6,7]

    def save_data(self,path,CFL):
        '''Store active variable in file. Path provided by solver'''
        outFileT = '/'.join([path,'_'.join([CFL,self.outFileT])])
        outFileA = '/'.join([path,'_'.join([CFL,self.outFileA])])
        with open(outFileT,'w') as out:
            if self.activeVariables == [0,1,2,3]:
                table = [self.ts,self.dts,self.avs,self.aves]
            elif self.activeVariables == [0,1,2,3,4,5,6,7]:
                table = [self.ts,self.dts,self.avs,self.aves,self.evs,self.eves,self.es,self.enrights]
            else:
                raise TypeError('Unknown format for self.activeVariables ',self.activeVariables)

            for i in range(len(table[0])):
                for j in range(len(table)):
                    out.write('%8E\t' % table[j][i])
                out.write('\n')
       
        with open(outFileA,'w') as out:
            for j in range(len(table)):
                avg = sum(array(table[j][:])*array(table[1][:]))/sum(array(table[1][:]))
                out.write('%8E\t' % avg)
            out.write('\n')

    def get_inflowBCS(self,t=0):
        '''Return a triplet w/ InflowBoundary object, the Dirichlet value and FacetFunction object.'''
        mesh = self.domain                   # get a domain from problem

        vx = self.vStrings[0]                # get velocity from problem
        vy = self.vStrings[1]
        v = Expression((vx,vy),t = t)
        VV = VectorFunctionSpace(mesh,'CG',1)   # space for velocity interpolation
        v = interpolate(v,VV)

        bcN = DirichletBC(VV,OutwardBoundaryVector(mesh),DomainBoundary())
        n = Function(VV)
        bcN.apply(n.vector())                    # n now has values of the outward surface normal on the boundary

        ib = InflowBoundary(n,v)                # InflowBoundary object, use to mark parts of boundary
                                                        
        bParts = MeshFunction('size_t',mesh,mesh.topology().dim()-1)
        bParts.set_all(0)                       # 0 denotes parts that will be used for surface integral
        ib.mark(bParts,1)                       # using ib, we define the parts of boundary for Derichlet as 1
        return (ib, self.ibcExpr, bParts)     #use 0,1 to prescribe Derichlet<---callable!, 2 forassemble

#<------------------------------------------METHODS TO BE DEFINED IN CHILDREN---------------------------------------------->

    def exact_volume(self):
        '''Return problem specific exact volume.'''
        raise NotImplementedError('Method exact_volume() should be defined in child classes.')

    def exact_solution(self,t=None):
        '''Return problem specific exact solution at time t.'''
        raise NotImplementedError('Method exact_solution() should be defined in child classes.')

